{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I1130 08:01:52.581158 139935788427072 file_utils.py:39] PyTorch version 0.4.1 available.\n"
     ]
    }
   ],
   "source": [
    "from create_inputs_utils import * \n",
    "import os\n",
    "import numpy as np\n",
    "import h5py\n",
    "import json\n",
    "from tqdm import tqdm\n",
    "from collections import Counter\n",
    "from random import seed, choice, sample\n",
    "import pickle\n",
    "import argparse\n",
    "import glob\n",
    "import logging\n",
    "import os\n",
    "import random\n",
    "import sys\n",
    "sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname('__file__'))))\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from tqdm import tqdm, trange\n",
    "import csv\n",
    "import logging\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,\n",
    "                              TensorDataset)\n",
    "from transformers import (WEIGHTS_NAME, BertConfig,\n",
    "                                  BertForSequenceClassification, BertTokenizer,\n",
    "                                  )\n",
    "\n",
    "dataset = 'coco'\n",
    "karpathy_json_path='../data/caption_datasets/dataset_coco.json'\n",
    "captions_per_image=5\n",
    "output_folder='../preprocessed_data'\n",
    "base_filename = 'preprocessed_' + dataset\n",
    "max_len=50\n",
    "\n",
    "train_image_captions = []\n",
    "val_image_captions = []\n",
    "test_image_captions = []\n",
    "\n",
    "train_image_det = []\n",
    "val_image_det = []\n",
    "test_image_det = []\n",
    "\n",
    "with open(karpathy_json_path, 'r') as j:\n",
    "    data = json.load(j)\n",
    "with open(os.path.join(output_folder,'train36_imgid2idx.pkl'), 'rb') as j:\n",
    "    train_data = pickle.load(j)       \n",
    "with open(os.path.join(output_folder,'val36_imgid2idx.pkl'), 'rb') as j:\n",
    "    val_data = pickle.load(j)\n",
    "    \n",
    "processor = CaptionProcessor()\n",
    "task_name = \"sst-2\"#임의 설정 (필요없는 값)\n",
    "output_mode = \"classification\" #임의 설정 (필요없는 값)\n",
    "model_name_or_path = \"bert-base-uncased\" \n",
    "max_seq_length = 50\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available()  else \"cpu\")\n",
    "label_list = processor.get_labels()\n",
    "num_labels = len(label_list)\n",
    "\n",
    "config_class = BertConfig\n",
    "model_class = BertForSequenceClassification\n",
    "tokenizer_class = BertTokenizer\n",
    "\n",
    "config = config_class.from_pretrained(model_name_or_path, num_labels=num_labels, finetuning_task = task_name)\n",
    "tokenizer = tokenizer_class.from_pretrained(model_name_or_path, do_lower_case = True)\n",
    "#model = model_class.from_pretrained(model_name_or_path, from_tf=bool('.ckpt' in model_name_or_path), config=config)\n",
    "\n",
    "#model.to(device)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_image_captions_with_len = []\n",
    "val_image_captions_with_len = []\n",
    "test_image_captions_with_len = []\n",
    "\n",
    "for img in data['images']:\n",
    "    captions = []\n",
    "    captions_with_len = []\n",
    "    \n",
    "    for caption in img['sentences']:\n",
    "        # Update word frequency\n",
    "        if len(caption['tokens']) <= max_len:\n",
    "            #captions_forlen : [['a','man','with','a','red','helmet'],['a','man',..]] 5개의 캡션씩 들어 있음\n",
    "            #['a','man','with','a','red','helmet'] -> \"a man with a red helmet\"\n",
    "            \n",
    "            #bowonko\n",
    "            caption_len = len(caption['tokens']) + 2\n",
    "            \n",
    "            caption_sen = \" \".join(caption['tokens'])\n",
    "             \n",
    "            captions.append(caption_sen)\n",
    "            \n",
    "            #bowonko\n",
    "            \n",
    "            cwl_item = (caption_sen,caption_len)\n",
    "            captions_with_len.append(cwl_item)\n",
    "            \n",
    "    #captions : ['sen1','sen2','sen3','sen4','sen5']\n",
    "    if len(captions) == 0:\n",
    "        print('절대 발생할 수 없음')\n",
    "        continue\n",
    "\n",
    "    #ID\n",
    "    image_id = img['filename'].split('_')[2]\n",
    "    image_id = int(image_id.lstrip(\"0\").split('.')[0])\n",
    "\n",
    "    #split은 train, val, test, restval 로 구성\n",
    "    if img['split'] in {'train', 'restval'}:\n",
    "        if img['filepath'] == 'train2014':\n",
    "            if image_id in train_data:\n",
    "                train_image_det.append((\"t\",train_data[image_id]))\n",
    "        else:\n",
    "            if image_id in val_data:\n",
    "                train_image_det.append((\"v\",val_data[image_id]))\n",
    "        #bowonko\n",
    "        train_image_captions.append(captions)\n",
    "        train_image_captions_with_len.append(captions_with_len)\n",
    "        \n",
    "        #train_image_captions은 2차원, element(captions) = ['sen1','sen2','sen3','sen4','sen5'], \n",
    "\n",
    "    elif img['split'] in {'val'}:\n",
    "        if image_id in val_data:\n",
    "            val_image_det.append((\"v\",val_data[image_id]))   \n",
    "        #bowonko    \n",
    "        val_image_captions.append(captions)\n",
    "        val_image_captions_with_len.append(captions_with_len)\n",
    "        \n",
    "        #val_image_captions은 2차원, element(captions) = ['sen1','sen2','sen3','sen4','sen5']\n",
    "    elif img['split'] in {'test'}:\n",
    "        if image_id in val_data:\n",
    "            test_image_det.append((\"v\",val_data[image_id])) \n",
    "        #bowonko\n",
    "        test_image_captions.append(captions)\n",
    "        test_image_captions_with_len.append(captions_with_len)\n",
    "        \n",
    "        #test_image_captions은 2차원, element(captions) = ['sen1','sen2','sen3','sen4','sen5']\n",
    "# Sanity check\n",
    "assert len(train_image_det) == len(train_image_captions)\n",
    "assert len(train_image_det) == len(train_image_captions_with_len)\n",
    "assert len(val_image_det) == len(val_image_captions)\n",
    "assert len(val_image_det) == len(val_image_captions_with_len)\n",
    "assert len(test_image_det) == len(test_image_captions)\n",
    "assert len(test_image_det) == len(test_image_captions_with_len)\n",
    "\n",
    "caption_size = 5\n",
    "\n",
    "#captions_forlen : [['a','man','with','a','red','helmet'],['a','man',..]]\n",
    "for impaths, imcaps, split in [(train_image_det, train_image_captions_with_len, 'TRAIN'),\n",
    "                                   (val_image_det, val_image_captions_with_len, 'VAL'),\n",
    "                                   (test_image_det, test_image_captions_with_len, 'TEST')]:\n",
    "    \n",
    "    #imcaps == image_captions_with_len > captions_with_len > (caption_sen,caption_len)\n",
    "    #cwl_item = (caption_sen,caption_len)\n",
    "    #captions_with_len.append(cwl_item)\n",
    "    \n",
    "    caplens = []\n",
    "    image_captions = []\n",
    "    enc_captions = []\n",
    "    \n",
    "    #captions에 (caption_sen,caption_len) 다 들어가게 변경했으므로 captions_sen에 caption_sen만 모아놓음\n",
    "    captions_sen = []\n",
    "    \n",
    "    for i, path in enumerate(tqdm(impaths)):\n",
    "        image_captions = []\n",
    "        \n",
    "        # Sample captions\n",
    "        #이런 일은 발생하지 않음\n",
    "        if len(imcaps[i]) < captions_per_image:#['sen1','sen2','sen3','sen4','sen5']\n",
    "            #이미지 당 caption이 5개보다 적으면 중복해서 구성\n",
    "            captions = imcaps[i] + [choice(imcaps[i]) for _ in range(captions_per_image - len(imcaps[i]))]\n",
    "        else: #모두 여기로 들어옴\n",
    "            captions = sample(imcaps[i], k=captions_per_image)\n",
    "\n",
    "        # Sanity check\n",
    "        assert len(captions) == captions_per_image\n",
    "                \n",
    "        for cap_sen,cap_len in captions:\n",
    "            caplens.append(cap_len)\n",
    "            captions_sen.append(cap_sen)\n",
    "        \n",
    "        image_captions.append(captions_sen)\n",
    "        \n",
    "    \n",
    "    \n",
    "    with open(os.path.join(output_folder, split + '_CAPLENS_' + base_filename + '.json'), 'w') as j:\n",
    "        json.dump(caplens, j)\n",
    "    \n",
    "    data_dir = output_folder\n",
    "\n",
    "\n",
    "    #목적 load_and_cache_example을 불러서 cache 파일 만들기 이후 두 번째 실행할 때는 load \n",
    "    if(split == \"TRAIN\"): \n",
    "        #TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids)\n",
    "        train_dataset = load_and_cache_examples(processor,tokenizer,max_seq_length,model_name_or_path,image_captions,data_dir,split)\n",
    "        caplens_train = []\n",
    "        batch_size = len(train_dataset)\n",
    "        \n",
    "        for b in range(batch_size):\n",
    "            train_caption = train_dataset[b][0]\n",
    "            cnt = 0\n",
    "            for i in range(len(train_caption)):\n",
    "                if(train_caption[i] == 0):\n",
    "                    break\n",
    "                else:\n",
    "                    cnt += 1\n",
    "            caplens_train.append(cnt)        \n",
    "                \n",
    "          \n",
    "        with open(os.path.join(output_folder, 'TRAIN' + '_CAPLENS_' + base_filename + '.json'), 'w') as j:\n",
    "            json.dump(caplens_train, j)\n",
    "        \n",
    "        \n",
    "    elif(split == \"VAL\"):\n",
    "        dev_dataset = load_and_cache_examples(processor,tokenizer,max_seq_length,model_name_or_path,image_captions,data_dir,split)\n",
    "        caplens_dev = []\n",
    "        batch_size = len(dev_dataset)\n",
    "        \n",
    "        for b in range(batch_size):\n",
    "            cnt = 0\n",
    "            dev_caption = dev_dataset[b][0]\n",
    "            for i in range(len(dev_caption)):\n",
    "                if(dev_caption[i] == 0):\n",
    "                    break\n",
    "                else:\n",
    "                    cnt += 1\n",
    "            caplens_dev.append(cnt)\n",
    "                                  \n",
    "        with open(os.path.join(output_folder, 'VAL' + '_CAPLENS_' + base_filename + '.json'), 'w') as j:\n",
    "            json.dump(caplens_dev, j)\n",
    "        \n",
    "    elif(split == \"TEST\"):\n",
    "        test_dataset = load_and_cache_examples(processor,tokenizer,max_seq_length,model_name_or_path,image_captions,data_dir,split)\n",
    "        caplens_test = []\n",
    "        batch_size = len(test_dataset)\n",
    "        for b in range(batch_size):\n",
    "            cnt = 0\n",
    "            test_caption = test_dataset[b][0]\n",
    "            for i in range(len(test_caption)):\n",
    "                if(test_caption[i] == 0):\n",
    "                    break\n",
    "                else:\n",
    "                    cnt += 1\n",
    "            caplens_test.append(cnt)\n",
    "                \n",
    "                    \n",
    "        with open(os.path.join(output_folder, 'TEST' + '_CAPLENS_' + base_filename + '.json'), 'w') as j:\n",
    "            json.dump(caplens_test, j)    \n",
    "\n",
    "            \n",
    "    \n",
    "# Save bottom up features indexing to JSON files\n",
    "with open(os.path.join(output_folder, 'TRAIN' + '_GENOME_DETS_' + base_filename + '.json'), 'w') as j:\n",
    "    json.dump(train_image_det, j)\n",
    "\n",
    "with open(os.path.join(output_folder, 'VAL' + '_GENOME_DETS_' + base_filename + '.json'), 'w') as j:\n",
    "    json.dump(val_image_det, j)\n",
    "\n",
    "with open(os.path.join(output_folder, 'TEST' + '_GENOME_DETS_' + base_filename + '.json'), 'w') as j:\n",
    "    json.dump(test_image_det, j)\n",
    "   \n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
